from common.arguments import get_common_args
from common.agent import Agent
import numpy as np
import gym

def evaluate(env, agent, args):
    #评估当前模型性能 
    total_reward = 0
    state = env.reset()
    for step in range(args.max_steps):
        action = agent.choose_action(state, exploration=False)
        next_state, reward, done, _ = env.step(action)
        total_reward += reward
        state = next_state                      
        # env.render()
        if done or step == args.max_steps-1:
            break
    return total_reward

def main():
    args = get_common_args()

    alg = "DQN"
    test_episodes = 100

    if alg == "DDPG":
        env_name = 'Pendulum-v1'
        model = "2112_211702_435episodes"
        env = gym.make(env_name)
        state_size = env.observation_space.shape[0]
        action_size = env.action_space.shape[0]
        action_bound = env.action_space.high[0]
        print("action_bound = ", action_bound)
        path = "model/{}/{}".format(alg, model)
    else:
        # for DQN DDQN RainbowDQN
        env_name = 'CartPole-v1'
        model = "2112_211410_650episodes"
        env = gym.make(env_name)
        state_size = env.observation_space.shape[0]
        action_size = env.action_space.n

        path = "model/{}/{}.pth".format(alg, model)
    
    #init the agent
    agent = Agent(alg, state_size, action_size, args)
    agent.load(path)

    reward_list = []
    for _ in range(test_episodes):
        reward = evaluate(env, agent, args)
        reward_list.append(reward)
    print("The {} average reward is: {}".format(alg, np.average(reward_list)))

if __name__ == '__main__':
    main()



